from operator import itemgetter
# langchain
from langchain_openai import ChatOpenAI 
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableWithMessageHistory,RunnablePassthrough,RunnableLambda
from langchain_community.chat_message_histories import SQLChatMessageHistory
from langchain.memory import ConversationSummaryBufferMemory
from langchain_core.messages import HumanMessage, AIMessage
from langchain.prompts import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder
)
from langchain_core.messages import BaseMessage
from servers.EmbeddingData import  search_similar_text
from servers.mysqlconnect import create_connection
from typing import List




class CustomConversationSummaryBufferMemory(ConversationSummaryBufferMemory):
    def __init__(self, *args, **kwargs):
        # 调用父类的构造函数
        super().__init__(*args, **kwargs)

    # 自定义 token 计数方法
    def get_num_tokens_from_messages(self, messages: List[BaseMessage]) -> int:
        """自定义 token 计数：根据消息内容估算 token 数量"""
        total_chars = 0
        for msg in messages:
            total_chars += len(str(msg.content))  # 获取消息内容的字符数
        return total_chars // 4  # 每个 token 估算为 4 个字符

    # 覆盖 prune 方法，使用自定义 token 计数方法
    def prune(self) -> None:
        """如果总 token 数量超出限制，则进行修剪"""
        buffer = self.chat_memory.messages
        curr_buffer_length = self.get_num_tokens_from_messages(buffer)  # 使用自定义的 token 计数方法
        if curr_buffer_length > self.max_token_limit:
            pruned_memory = []
            while curr_buffer_length > self.max_token_limit:
                pruned_memory.append(buffer.pop(0))
                curr_buffer_length = self.get_num_tokens_from_messages(buffer)
            self.moving_summary_buffer = self.predict_new_summary(
                pruned_memory,
                self.moving_summary_buffer,
            )


# LLM相关功能
def llm_stream(text,subjectid:int):
    connection = create_connection()
    connection.ping(reconnect=True)
    # 创建模型
    chat_model = ChatOpenAI(
        api_key="sk-b07d91eefa82450d90448d549e87d4b7",
        base_url="https://api.deepseek.com/v1",
        model="deepseek-chat",
        timeout=30,
        max_retries=3
    )
    # 提示词模板
    # template="""
    # 请根据以下内容回答问题:
    # {context}
    # 用户的问题是:{question}
    # """
    
    # retrieve_context 修改：格式化为纯文本块（全文）
    def retrieve_context(question: str) -> str:
        results = search_similar_text(question, top_k=3)  # list[dict]
        if not results:
            return "无相关知识信息。"
        context = "知识信息：\n"
        for i, res in enumerate(results, 1):
            context += f"[文档{i}]: {res['text']}\n\n"  # 全文，无截取
        return context  # ~6k 字符，安全
    
    
    # 系统的消息,用于插入检索后的相关文本
    system_message = SystemMessagePromptTemplate.from_template(
        """根据以下已知信息和历史对话回答用户问题。如果历史中已有相关信息，请参考它。
        已知信息：{context}

        **重要：使用Markdown格式来组织你的回答，包括：**
        - 使用标题（#、##、###）来组织内容结构
        - 使用粗体（**粗体**）强调重要概念
        - 使用代码块（```代码```）展示技术示例
        - 使用列表（- 或 1.）来列举要点
        - 使用表格来对比不同方案
        - 使用链接格式 [文本](URL) 来引用资源
        """
    )
    # 创建一个人类消息，用于接收用户的输入
    human_message = HumanMessagePromptTemplate.from_template("用户问题：{question}")
    prompt = ChatPromptTemplate.from_messages([system_message, MessagesPlaceholder(variable_name="history"), human_message])
    # 输出解析器
    out_put_parser = StrOutputParser()
    

    
    # 多轮记忆：用自定义类 (修复 token 错误) (总结旧 + 保留最近 3 轮)
    memory = CustomConversationSummaryBufferMemory(
    llm=chat_model,  # 使用的 LLM
    max_token_limit=3000,  # 最大 token 限制
    return_messages=True,
    memory_key="history",  # 记忆的 key
)
    
    # 加载/保存历史到 SQL (自定义，结合 Memory)
    def load_save_history(session_id: str):
        # 从 SQL 加载旧历史 (chatcontent 表)
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT answer, question FROM chatcontent WHERE subject_id = %s ORDER BY created_at",
                (session_id,)
            )
            rows = cursor.fetchall()
            history_dicts = []
            for row in rows:
                history_dicts.append({"role": "human", "content": str(row['question'])})
                history_dicts.append({"role": "ai", "content": str(row['answer'])})
        
        history_msgs = []
        for d in history_dicts[-6:]:  # 最近 3 轮 (user+ai)
            try:
                if d["role"] == "human":
                    history_msgs.append(HumanMessage(content=d["content"]))
                else:
                    history_msgs.append(AIMessage(content=d["content"]))
            except Exception as e:
                print(f"Message conversion error: {e}, skipping.")
                continue
        # 总结 + 注入 memory
        memory.chat_memory.add_messages(history_msgs)  
        memory.save_context({"input": text}, {"output": "temp"})  # 保存当前 (手动)
        
        return memory.chat_memory.messages  # 返回注入后的历史

    # 基础链  {"question": text} -> search_similar_text检索相识知识 -> 将检索结果和问题插入到 prompt 中 -> 模型 -> 输出解析器 -> 输出
    chains=(
        # RunnablePassthrough.assign：从输入中“借道”传递所有字段，同时动态添加context键：用itemgetter("question")提取问题
        RunnablePassthrough.assign(  # 动态添加 context
            context=itemgetter("question") | RunnableLambda(retrieve_context)
        )
        |prompt
        |chat_model
        |out_put_parser
    )
    
    # 加载历史
    history = load_save_history(str(subjectid))
    
    
    # 调用chain链
    config = {"configurable": {"session_id": subjectid}}
    try:
        result = chains.stream(
            {"question": text, "history": history},  # 注入历史
            config=config
        )
        return result
    except Exception as e:
        # 如果stream调用失败，返回一个生成器输出错误信息
        def error_generator():
            yield f"错误: {str(e)}"
        return error_generator()


# 主题生成
def llm(text):
    # 创建模型
    try:
        chat_model = ChatOpenAI(
            api_key="sk-b07d91eefa82450d90448d549e87d4b7",
            base_url="https://api.deepseek.com/v1",
            model="deepseek-chat",
            timeout=30,
            max_retries=3
        )
        # 提示词模板
        template="""
        请根据以下内容生成会话主题:
        {text}
        返回的结果只有主题，不要返回其他任何内容。
        """
        system_message = SystemMessagePromptTemplate.from_template(template)
        prompt = ChatPromptTemplate.from_messages([system_message])
        out_put_parser = StrOutputParser()
        
        # 正确的调用链构建
        chains = prompt | chat_model | out_put_parser
        
        print(f"尝试生成主题，输入文本: {text}")
        
        # 正确的参数传递和返回值处理
        try:
            result = chains.invoke({"text": text})
            print(f"主题生成成功，结果: {result}")
            return result
        except Exception as e:
            print(f"主题生成失败 (API调用错误): {str(e)}")
            raise Exception(f"API调用错误: {str(e)}")
    except Exception as e:
        print(f"主题生成失败 (初始化错误): {str(e)}")
        raise Exception(f"模型初始化错误: {str(e)}")